# coding=utf-8
# 用opencv实现兑换站识别
import cv2
import numpy as np
import pyrealsense2 as rs
import time
import threading
#ros发布的部分
# import rospy
# from geometry_msgs.msg import Point, Vector3

info = [[0,0,0],[0,0,0],[0,0,0]]

# 求取图像上一个点的三维坐标
def get_3d_point(depth_frame, xy):
    x, y = int(xy[0]), int(xy[1])
    # 获取像素点的深度值
    depth = depth_frame.get_distance(x, y)
    # 如果是噪点，则用四周不为0的点的深度值的平均值代替
    if depth < 0.01:
        count = 0
        for i in range(-3, 4):
            for j in range(-3, 4):
                # 防止越界
                x_limit = max(min(x+i, W-1), 0)
                y_limit = max(min(y+j, H-1), 0)
                if depth_frame.get_distance(x_limit, y_limit) > 0.01:
                    depth += depth_frame.get_distance(x_limit, y_limit)
                    count += 1
        if count != 0:
            depth = depth / count
    if depth == 0:
        print("depth is 0")
    # intelrealsense参数
    intrinsics = depth_frame.profile.as_video_stream_profile().intrinsics
    # 像素坐标系转换为相机坐标系
    camera_coordinate = rs.rs2_deproject_pixel_to_point(intrinsics, [x, y], depth)
    return camera_coordinate

# 三点坐标 -> 正方体前表面坐标，正方体三个互相垂直面的单位法向量
def getCenter(p1, p2, p3, roll_45=False, direction="right"):
    p1 = np.array(p1)
    p2 = np.array(p2)
    p3 = np.array(p3)
    # 化p1,p3在同一条对角线上
    p1p2 = np.linalg.norm(p1 - p2)
    p1p3 = np.linalg.norm(p1 - p3)
    p2p3 = np.linalg.norm(p2 - p3)
    if(p1p2 > p1p3 and p1p2 > p2p3):
        p3, p2 = p2, p3
    elif(p2p3 > p1p2 and p2p3 > p1p3):
        p1, p2 = p2, p1
    # 化p1为左上角点，p3为右下角点
    # if p1[0] + p1[1] > p3[0] + p3[1]:
    #     p1, p3 = p3, p1
    # p1p2向量(一个法向量)
    vec2 = np.array(p2) - np.array(p1)
    # p2p3向量(另一个法向量)
    vec3 = np.array(p3) - np.array(p2)
    # 正面法向量
    vec1 = np.cross(vec2, vec3)
    # 法向量单位化
    vec_norm1 = np.linalg.norm(vec1)
    unit_vec1 = vec1 / (vec_norm1 if vec_norm1 != 0 else 1)
    vec_norm2 = np.linalg.norm(vec2)
    unit_vec2 = vec2 / (vec_norm2 if vec_norm2 != 0 else 1)
    vec_norm3 = np.linalg.norm(vec3)
    unit_vec3 = vec3 / (vec_norm3 if vec_norm3 != 0 else 1)

    
    # 旋转45度
    if roll_45:
        unit_vec2 = (unit_vec2 + unit_vec3) / 2 ** 0.5
        unit_vec3 = (unit_vec2 - unit_vec3) / 2 ** 0.5
        unit_vec2 = unit_vec2 / np.linalg.norm(unit_vec2)
        unit_vec3 = unit_vec3 / np.linalg.norm(unit_vec3)
    
    # z轴法向量指向正方向
    if(unit_vec1[2] < 0):
        unit_vec1 = -unit_vec1
    # 平面中心
    front_center = (p1 + p3) / 2
    # 固定棱长(m)
    gold_a = 0.24
    # 计算棱长
    # gold_a = np.linalg.norm(p1 - p2)
    # 中心坐标
    center = front_center + gold_a / 2 * unit_vec1
    # t=体心坐标，z轴方向单位法向量，另外两个方向单位向量
    # return center, unit_vec1, unit_vec2, unit_vec3
    vec_z = unit_vec1
    if unit_vec2[0] < 0:
        unit_vec2 = -unit_vec2
    if unit_vec3[0] < 0:
        unit_vec3 = -unit_vec3
    if np.dot(unit_vec2, np.array([1, 0, 0])) > np.dot(unit_vec3, np.array([1, 0, 0])):
        vec_x = unit_vec2
        vec_y = unit_vec3
    else:
        vec_x = unit_vec3
        vec_y = unit_vec2
    if vec_x[0] < 0:
        vec_x = -vec_x
    if vec_y[1] < 0:
        vec_y = -vec_y
    # 兑换站中心、前方向、下方向、右方向
    # return center, vec_z, vec_x, vec_y
    # return center, vec_z, vec_x, vec_y
    if roll_45 and direction == "right":
        vec_z, vec_x, vec_y = -vec_y, vec_x, vec_z
    elif roll_45 and direction == "left":
        vec_z, vec_x, vec_y = vec_y, vec_x, vec_z
    return front_center, vec_z, vec_x, vec_y

# 输出结果
def printResult(depth_frame, points, roll_45 = False, direction="right"):
    global info
    # 绘制points中的点
    for j in range(len(points)):
        cv2.circle(color_image, (points[j][0], points[j][1]), 4, (255, 255, 255), -1)
    # 读取points中的点的三维坐标
    points_3d = []
    for i in range(3):
        dep = get_3d_point(depth_frame, points[i])
        points_3d.append(dep)
    # 输出结果
    if len(points_3d) >= 3:
        output = getCenter(points_3d[0], points_3d[1], points_3d[2], roll_45, direction)
        print("center: ", output[0])
        print("front_vector: ", output[1])
        print("right_vector: ", output[2])

        # lock.acquire()
        info = [output[0], output[1], output[2]]
        # print("lock")
        # lock.release()

    print("")

# 两点距离
def getDistance(p1, p2):
    p1 = np.array(p1)
    p2 = np.array(p2)
    return np.linalg.norm(p1 - p2)

def points_in_contour(contour):
    inside = []
    rect = cv2.boundingRect(contour)
    # 取boundingrect内的所有像素
    for m in range(rect[0], rect[0] + rect[2]):
        for n in range(rect[1], rect[1] + rect[3]):
            # 判断点是否在旋转矩形内
            if cv2.pointPolygonTest(contour, (m, n), False) > 0:
                inside.append([m, n])
    return inside

# L形轮廓的两边交点
def contour_point(contour):
            points = []            
            # 轮廓中最远的两个点
            max_l = 0
            max_idx = [0, 0]
            for j in range(len(contour)):
                for k in range(j + 1, len(contour)):
                    l = (contour[j][0][0] - contour[k][0][0]) ** 2 + \
                        (contour[j][0][1] - contour[k][0][1]) ** 2
                    if l > max_l:
                        max_l = l
                        max_idx = [j, k]
            points.append(contour[max_idx[0]][0])
            points.append(contour[max_idx[1]][0])
            # 轮廓中距离这两个点距离之和最远的点
            max_l = 0
            max_idx = 0
            for j in range(len(contour)):
                l = (contour[j][0][0] - contour[max_idx][0][0]) ** 2 + \
                    (contour[j][0][1] - contour[max_idx][0][1]) ** 2
                l = getDistance(points[0], contour[j][0]) + \
                    getDistance(points[1], contour[j][0])
                if l > max_l:  
                    max_l = l
                    max_idx = j
            points.append(contour[max_idx][0])
            # points中的直角点
            point_90 = points[0]
            max_dis = 0
            for j in range(3):
                for k in range(j + 1, 3):
                    center_2 = (np.array(points[j]) + np.array(points[k])) / 2
                    center_2 = center_2.astype(np.int32)
                    # cv2.circle(color_image, (center_2[0], center_2[1]), 4, (0, 255, 255), -1)
                    min_dis = 1000000
                    inside = points_in_contour(contour)
                    for p in inside:
                        dis = (p[0] - center_2[0]) ** 2 + \
                              (p[1] - center_2[1]) ** 2
                        if dis < min_dis:
                            min_dis = dis
                    x = int(center_2[0])
                    y = int(center_2[1])
                    if cv2.pointPolygonTest(contour, (x, y), False) > 0:
                        min_dis = 0
                    print(min_dis)
                    if min_dis > max_dis:
                        max_dis = min_dis
                        point_90 = points[3 - j - k]
            cv2.circle(color_image, (point_90[0], point_90[1]), 4, (255, 255, 0), -1)
            # 斜边中点
            center = np.array([0, 0])
            for j in range(3):
                if points[j][0] != point_90[0] or points[j][1] != point_90[1]:
                    center += np.array(points[j])
            center = center / 2 # 除以2后为浮点数
            cv2.circle(color_image, (int(center[0]), int(center[1])), 4, (255, 0, 255), -1)

            # 轮廓分割(取半部分)
            # 可以画线来分
            center_line = np.array(center) - np.array(point_90)
            count = len(contour)
            right = []
            left = []
            points_inside = points_in_contour(contour)
            for j in range(len(points_inside)):
                vector = np.array(points_inside[j]) - np.array(point_90)
                if np.cross(center_line, vector) > 0:
                    right.append(points_inside[j])
                else:
                    left.append(points_inside[j])
            right = np.array(right)
            left = np.array(left)
            right_rect = cv2.minAreaRect(right)
            left_rect = cv2.minAreaRect(left)
            box1 = cv2.boxPoints(right_rect)
            box2 = cv2.boxPoints(left_rect)
            # 取旋转矩形内的点用来拟合直线
            # 右
            right = []
            rect = cv2.boundingRect(box1)
            # 取boundingrect内的所有像素
            for j in range(rect[0], rect[0] + rect[2]):
                for k in range(rect[1], rect[1] + rect[3]):
                    # 判断点是否在旋转矩形内
                    if cv2.pointPolygonTest(box1, (j, k), False) > 0:
                        right.append([j, k])
            # 将right转换为int
            right = np.array(right)
            right = right.astype(int)
            right_line = cv2.fitLine(right, cv2.DIST_L2, 0, 0.01, 0.01)
            cv2.line(color_image, (int(right_line[2] - right_line[0] * 100), int(right_line[3] - right_line[1] * 100)), (int(right_line[2] + right_line[0] * 100), int(right_line[3] + right_line[1] * 100)), (0, 255, 255), 2)
            right_vec = np.array([right_line[0][0] , right_line[1][0]])
            right_vec = right_vec / (np.linalg.norm(right_vec) if np.linalg.norm(right_vec) != 0 else 1)
            if np.cross(center_line, right_vec) < 0:
                right_vec = -right_vec
            # 左
            left = []
            rect = cv2.boundingRect(box2)
            for j in range(rect[0], rect[0] + rect[2]):
                for k in range(rect[1], rect[1] + rect[3]):
                    # 判断点是否在旋转矩形内
                    if cv2.pointPolygonTest(box2, (j, k), False) > 0:
                        left.append([j, k])
            left_line = cv2.fitLine(np.array(left), cv2.DIST_L2, 0, 0.01, 0.01)
            cv2.line(color_image, (int(left_line[2] - left_line[0] * 100), int(left_line[3] - left_line[1] * 100)), (int(left_line[2] + left_line[0] * 100), int(left_line[3] + left_line[1] * 100)), (0, 255, 255), 2)
            left_vec = np.array([left_line[0][0] , left_line[1][0]])
            left_vec = left_vec / (np.linalg.norm(left_vec) if np.linalg.norm(left_vec) != 0 else 1)
            if np.cross(center_line, left_vec) > 0:
                left_vec = -left_vec
            # 求right_line和left_line的交点
            a1, b1, c1 = right_line[1], -right_line[0], -right_line[1] * right_line[2] + right_line[0] * right_line[3]
            a2, b2, c2 = left_line[1], -left_line[0], -left_line[1] * left_line[2] + left_line[0] * left_line[3]
            x = (b1 * c2 - b2 * c1) /( (a1 * b2 - a2 * b1) if (a1 * b2 - a2 * b1) != 0 else 1)
            y = (a2 * c1 - a1 * c2) / ((a1 * b2 - a2 * b1) if (a1 * b2 - a2 * b1) != 0 else 1)
            return int(x), int(y)

# ROS_publish_topic,main中需要调用此函数
# def publish_geometry():

    # rospy.init_node('publish_geometry', anonymous=True,  disable_signals=True)
#     rate = rospy.Rate(10) # 10hz

#     # 创建四个消息发布者
#     center_pub = rospy.Publisher('center', Point, queue_size=10)
#     unit_vec1_pub = rospy.Publisher('unit_vec1', Point, queue_size=10)
#     unit_vec2_pub = rospy.Publisher('unit_vec2', Point, queue_size=10)
    
#     while (not rospy.is_shutdown()) and info[0][0] != 1000000:
#         # lock.acquire()
#         center, unit_vec1, unit_vec2  = info[0], info[1], info[2]
#         # lock.release
#         print("send")
#         # 发布center
#         center_msg = Point()
#         center_msg.x = -center[0]
#         center_msg.y = -center[1]
#         center_msg.z = center[2]
#         center_pub.publish(center_msg)

#         # 发布unit_vec1
#         unit_vec1_msg = Point()
#         unit_vec1_msg.x = -unit_vec1[0]
#         unit_vec1_msg.y = -unit_vec1[1]
#         unit_vec1_msg.z = unit_vec1[2]
#         unit_vec1_pub.publish(unit_vec1_msg)

#         # 发布unit_vec2
#         unit_vec2_msg =  Point()
#         unit_vec2_msg.x = -unit_vec2[0]
#         unit_vec2_msg.y = -unit_vec2[1]
#         unit_vec2_msg.z = unit_vec2[2]
#         unit_vec2_pub.publish(unit_vec2_msg)

#         rate.sleep()
#         print(info)

if __name__ == '__main__':

    # 参数设置
    W, H = 640, 480 # 图像宽高
    # W, H = 1280, 720 # 图像宽高
    FPS = 30 # 帧率
    MAX_LENGTH = 1000 # 区分正面和侧面轮廓所使用的空间长度
    RED = (0, 0, 255) # 红色
    BLUE = (255, 0, 0) # 蓝色
    GREEN = (0, 255, 0) # 绿色
    

    # 设置hsv阈值
    lower_blue = np.array([100, 43, 46])
    upper_blue = np.array([124, 255, 255])
    lower_red = np.array([0, 33, 40])
    # lower_red = np.array([0, 33, 10])
    upper_red = np.array([10, 255, 255])
    # 创建一个管道
    pipeline = rs.pipeline() 
    # 配置要流式传输的管道
    config = rs.config()
    config.enable_stream(rs.stream.color, W, H, rs.format.bgr8, 30) # 8位rgb图像
    config.enable_stream(rs.stream.depth, W, H, rs.format.z16, 30) # 16位深度图像
    # 左右双目
    config.enable_stream(rs.stream.infrared, 1, W, H, rs.format.y8, 30) # 8位左红外图像
    config.enable_stream(rs.stream.infrared, 2, W, H, rs.format.y8, 30) # 8位右红外图像
    align = rs.align(rs.stream.color)
    # 视野大，但有残缺
    # align = rs.align(rs.stream.depth)
    # 开始传输帧
    profile = pipeline.start(config)
    # 更改rgb相机的曝光值
    color_sensor = profile.get_device().first_color_sensor()
    depth_sensor = profile.get_device().first_depth_sensor()
    # 关闭自动曝光
    color_sensor.set_option(rs.option.enable_auto_exposure, 0)
    depth_sensor.set_option(rs.option.enable_auto_exposure, 0)
    

    # 设置曝光值
    color_sensor.set_option(rs.option.exposure, 25)
    depth_sensor.set_option(rs.option.exposure, 15000)
    # 
    depth_sensor = profile.get_device().first_depth_sensor()
    # 
    depth_scale = depth_sensor.get_depth_scale()

    sender_running = True
    # sender =threading.Thread(target=getDistance)
    # sender =threading.Thread(target=publish_geometry)
    # sender.start()
    # lock = threading.Lock()

    # 对视频进行循环处理
    while True:
        # 等待开启通道
        frames = pipeline.wait_for_frames()
        # 将深度图和RGB图对齐
        frames = align.process(frames)
        # 获得RGB图像
        color_frame = frames.get_color_frame()
        # 未获得RGB图像
        if not color_frame:
            continue
        # 获得深度图像
        depth_frame = frames.get_depth_frame()
        # 把图像像素转化为数组
        color_image = np.asanyarray(color_frame.get_data())
        depth_image = np.asanyarray(depth_frame.get_data())
        # 深度图像转化为伪彩色图像(alpha数值越大对比度色差越大)
        depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.21), cv2.COLORMAP_JET)
        # 转换到HSV (颜色空间转换) (色调H，饱和度S，亮度V)
        hsv = cv2.cvtColor(color_image, cv2.COLOR_BGR2HSV)
        # 根据阈值构建掩模 (将在两个阈值内的像素值设置为白色，而不在阈值区间内的像素值设置为黑色)
        binary = cv2.inRange(hsv, lower_red, upper_red)
        # 形态学处理
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        # binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
        # binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        # 放慢视频播放速度
        time.sleep(0.1)
        

        # 寻找轮廓
        contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        # 正面轮廓
        front_contours = []
        # 侧面轮廓
        side_contours = []

        # 求最大轮廓面积
        max_area = 0
        for i in range(len(contours)):
            area = cv2.contourArea(contours[i])
            if area > max_area:
                max_area = area
        
        # 遍历轮廓分类
        for i in range(len(contours)):
            # 筛去面积过小的轮廓
            area = cv2.contourArea(contours[i])
            if area < max_area / 7:
                continue

            # 筛去面积占外接矩形比例过大的轮廓
            rect =  cv2.minAreaRect(contours[i])
            rect_area = rect[1][0] * rect[1][1]
            if area > 0.9 * rect_area:
                continue

            # 轮廓中最远的两个点
            max_l = 0
            max_idx = [0, 0]
            for j in range(len(contours[i])):
                for k in range(j + 1, len(contours[i])):
                    l = (contours[i][j][0][0] - contours[i][k][0][0]) ** 2 + \
                        (contours[i][j][0][1] - contours[i][k][0][1]) ** 2
                    if l > max_l:
                        max_l = l
                        max_idx = [j, k]
            # 根据轮廓中最远的两个点的空间距离区分轮廓类别
            point_3d1 = get_3d_point(depth_frame, contours[i][max_idx[0]][0])
            point_3d2 = get_3d_point(depth_frame, contours[i][max_idx[1]][0])
            distance = getDistance(point_3d1, point_3d2)
            # 正面轮廓
            if(distance < 0.10):
                front_contours.append(contours[i])
            # 侧面轮廓
            elif distance > 0.10 and distance < 0.30:
                side_contours.append(contours[i])

            else:
                front_contours.append(contours[i])

        # for i in range(len(front_contours)):
        #     rect = cv2.minAreaRect(front_contours[i])
        #     box = cv2.boxPoints(rect)
        #     box = np.int0(box)
        #     cv2.drawContours(color_image, [box], 0, (255, 255, 0), 2)


        # 保留最大的3个正面轮廓
        while len(front_contours) > 3:
            min_area = 100000
            min_idx = 0
            for i in range(len(front_contours)):
                rect =  cv2.minAreaRect(front_contours[i])
                area = rect[1][0] * rect[1][1]
                if area < min_area:
                    min_area = area
                    min_idx = i
            front_contours.pop(min_idx)

         # 绘制轮廓外接矩形
        # for i in range(len(front_contours)):
        #     rect = cv2.minAreaRect(front_contours[i])
        #     box = cv2.boxPoints(rect)
        #     box = np.int0(box)
        #     cv2.drawContours(color_image, [box], 0, BLUE, 2)
        for i in range(len(side_contours)):
            rect = cv2.minAreaRect(side_contours[i])            
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            cv2.drawContours(color_image, [box], 0, GREEN, 2)

        # 正面轮廓处理
        if len(front_contours) == 4 or len(front_contours) == 3:

            rect_center = []
            for i in range(len(front_contours)):
                rect = cv2.minAreaRect(front_contours[i])
                rect_center.append(rect[0])

            # 3个矩形中心的中心点
            x = (rect_center[0][0] + rect_center[1][0] + rect_center[2][0]) / 3
            y = (rect_center[0][1] + rect_center[1][1] + rect_center[2][1]) / 3

            # 每个轮廓中距离中心点最远的点
            center = []   
            for i in range(len(front_contours)):
                max_l = 0
                max_idx = 0
                for j in range(len(front_contours[i])):
                    l = (front_contours[i][j][0][0] - x) ** 2 + \
                        (front_contours[i][j][0][1] - y) ** 2
                    if l > max_l:
                        max_l = l
                        max_idx = j
                center.append(front_contours[i][max_idx][0])

            # 取空间最长边中点为正面的中心点
            center_3d = []
            for i in range(3):
                center_3d.append(get_3d_point(depth_frame, center[i]))

            # 空间最长边
            max_length = 0
            max_index = [0, 1]
            for i in range(3):
                for j in range(i + 1, 3):
                    if getDistance(center_3d[i], center_3d[j]) > max_length:
                        max_length = getDistance(center_3d[i], center_3d[j])
                        max_index = [i, j]

            # 最小二乘法求L形两边的交点
            cross1 = contour_point(front_contours[max_index[0]])
            cross2 = contour_point(front_contours[max_index[1]])
            cross3 = contour_point(front_contours[3 - max_index[0] - max_index[1]])


            # 中心点坐标
            x = (cross1[0] + cross2[0]) / 2
            y = (cross1[1] + cross2[1]) / 2
            
            # 绘制中心点
            cv2.circle(color_image, (int(x), int(y)), 5, RED, -1)
            
            points = [cross1, cross2, cross3]

            out_of_range = False
            for p in points:
                if p[0] < 0 or p[0] > 640 or p[1] < 0 or p[1] > 480:
                    print("cross out of range")
                    out_of_range = True

            if out_of_range:
                continue
            # 输出结果
            printResult(depth_frame, points)

        # 侧面轮廓处理
        elif len(side_contours) > 0:
            for i in range(len(side_contours)):
                big_rect = cv2.minAreaRect(side_contours[i])
                box = cv2.boxPoints(big_rect)
                box = np.int0(box)
                center = (int(big_rect[0][0]), int(big_rect[0][1]))
                points = []
                # 最下点
                max_y = 0
                max_index = 0
                for j in range(len(side_contours[i])):
                    if side_contours[i][j][0][1] > max_y:
                        max_y = side_contours[i][j][0][1]
                        max_index = j
                points.append(side_contours[i][max_index][0])
                # 距下点最远的点
                max_l = 0
                max_index = 0
                for j in range(len(side_contours[i])):
                    l = getDistance(side_contours[i][j][0], points[0])
                    if l > max_l:
                        max_l = l
                        max_index = j
                points.append(side_contours[i][max_index][0])
                # 距两点距离之和最远的点
                max_l = 0
                max_index = 0
                for j in range(len(side_contours[i])):
                    l = getDistance(side_contours[i][j][0], points[0]) + getDistance(side_contours[i][j][0], points[1])
                    if l > max_l:
                        max_l = l
                        max_index = j
                points.append(side_contours[i][max_index][0])
                # 直角点
                max_l = 0
                max_index = 0
                for j in range(len(points)):
                    for k in range(j + 1, len(points)):
                        l = getDistance(get_3d_point(depth_frame, points[j]), get_3d_point(depth_frame, points[k]))
                        if l > max_l:
                            max_l = l
                            max_index = 3 - j - k
                point = points[max_index]
                # big_rect的中心点
                center = (int(big_rect[0][0]), int(big_rect[0][1]))
                # 判断方向
                direction = "right" if point[0] < center[0] else "left"
                print(direction)
                # 输出结果
                printResult(depth_frame, points, True, direction)

        # 识别失败
        if len(front_contours) > 4:
            print("too many contours")
        elif len(front_contours) < 3 and len(side_contours) == 0:
            print("too few contours")

        # 显示图像
        cv2.imshow('color_image', color_image)
        cv2.imshow('depth_image', depth_colormap)
        cv2.imshow('binary', binary)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    info[0][0] = 1000000

    # 杀死sender子线程
    # sender.kill()

    cv2.destroyAllWindows()
